# video_inference_rtsp_vehicle_segmentation.py
import cv2
import numpy as np
from openvino.runtime import Core
import os
import time
import warnings
import logging
from ultralytics import YOLO

# -----------------------------
# 配置路径（请根据你的实际路径修改！）
# -----------------------------
YOLOV8_MODEL_PATH = r"models/yolov8n-seg.pt"  # YOLOv8n-seg模型路径
OUTPUT_FOLDER = r"detect"  # 输出结果文件夹
MASK_FOLDER = r"masks"  # 车辆掩码图片文件夹
ORIGINAL_FOLDER = r"original"  # 原始图片文件夹

# RTSP相机配置
RTSP_URLS = [
    "rtsp://admin:wkkj1234@192.168.200.74:554/Steaming/Channels/1",  # 示例RTSP URL，根据实际修改
]

# 创建输出目录
os.makedirs(OUTPUT_FOLDER, exist_ok=True)
os.makedirs(MASK_FOLDER, exist_ok=True)
os.makedirs(ORIGINAL_FOLDER, exist_ok=True)  # 创建原始图片文件夹

# -----------------------------
# 配置参数
# -----------------------------
DETECTION_INTERVAL = 3  # 视频检测间隔帧数
YOLO_CONFIDENCE = 0.3  # YOLO检测置信度阈值
YOLO_IOU = 0.5  # YOLO IOU阈值
SAVE_MASKS = True  # 是否保存车辆掩码图片
SAVE_ORIGINAL = True  # 是否保存原始图片（仅在检测到车辆时）
MASK_ALPHA = 0.5  # 掩码透明度（提高透明度使掩码更明显）
RTSP_TIMEOUT = 10  # RTSP连接超时时间（秒）
RTSP_RECONNECT_INTERVAL = 5  # RTSP重连间隔（秒）

# 设置环境变量
os.environ['GLOG_minloglevel'] = '2'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
warnings.filterwarnings('ignore')

# -----------------------------
# 1. 加载YOLOv8n-seg实例分割模型
# -----------------------------
print("Loading YOLOv8n-seg instance segmentation model...")
try:
    yolo_model = YOLO(YOLOV8_MODEL_PATH)
    print("✅ YOLOv8n-seg模型加载成功")
except Exception as e:
    print(f"❌ YOLOv8n-seg模型加载失败: {e}")
    print("请确保已安装ultralytics: pip install ultralytics")
    exit(1)

# -----------------------------
# 2. 车辆类别定义和颜色映射
# -----------------------------
VEHICLE_CLASSES = {
    2: 'car',  # 汽车
    3: 'motorcycle',  # 摩托车
    5: 'bus',  # 公交车
    7: 'truck',  # 卡车
}

VEHICLE_COLORS = {
    2: (0, 255, 0),  # 绿色 - 汽车
    3: (255, 255, 0),  # 青色 - 摩托车
    5: (255, 0, 255),  # 紫色 - 公交车
    7: (0, 165, 255),  # 橙色 - 卡车
}

# -----------------------------
# 3. 使用YOLOv8进行车辆检测和实例分割
# -----------------------------
def detect_vehicles_yolov8(image):
    """
    使用YOLOv8进行车辆检测和实例分割
    返回: 检测结果列表，每个元素为 [x_min, y_min, x_max, y_max, confidence, class_id, mask]
    """
    try:
        # 使用YOLOv8进行推理
        results = yolo_model(image, conf=YOLO_CONFIDENCE, iou=YOLO_IOU, classes=list(VEHICLE_CLASSES.keys()))
        detections = []

        for result in results:
            if hasattr(result, 'boxes') and result.boxes is not None:
                boxes = result.boxes
                masks = result.masks if hasattr(result, 'masks') and result.masks is not None else None

                for i, box in enumerate(boxes):
                    # 获取边界框坐标
                    x_min, y_min, x_max, y_max = box.xyxy[0].cpu().numpy()
                    confidence = box.conf[0].cpu().numpy()
                    class_id = int(box.cls[0].cpu().numpy())

                    # 获取分割掩码
                    mask_data = None
                    if masks is not None and i < len(masks.data):
                        # 获取当前检测框对应的掩码
                        mask_tensor = masks.data[i]
                        if mask_tensor is not None:
                            # 将掩码转换为numpy数组并调整到原图尺寸
                            orig_h, orig_w = image.shape[:2]
                            mask_resized = mask_tensor.cpu().numpy()

                            # 确保掩码是2D数组
                            if len(mask_resized.shape) == 3:
                                mask_resized = mask_resized[0]  # 取第一个通道

                            # 调整掩码尺寸到原图大小
                            if mask_resized.shape[0] != orig_h or mask_resized.shape[1] != orig_w:
                                mask_data = cv2.resize(mask_resized, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
                            else:
                                mask_data = mask_resized

                    detections.append([
                        int(x_min), int(y_min), int(x_max), int(y_max),
                        float(confidence), class_id, mask_data
                    ])

        print(f"    🔍 检测到 {len(detections)} 个车辆目标")
        return detections

    except Exception as e:
        print(f"    ❌ YOLOv8检测出错: {e}")
        return []

# -----------------------------
# 4. 保存原始图片函数（仅在检测到车辆时调用）
# -----------------------------
def save_original_image(image, frame_info="", vehicle_count=0):
    """
    保存原始图片，不带任何检测信息
    仅在检测到车辆时保存
    """
    if not SAVE_ORIGINAL or vehicle_count == 0:
        return None

    timestamp = time.strftime("%Y%m%d_%H%M%S")

    if "RTSP" in frame_info:
        # RTSP流帧
        frame_num = frame_info.split("_")[-1] if "_" in frame_info else "unknown"
        original_filename = f"original_rtsp_frame_{frame_num}_vehicles{vehicle_count}_{timestamp}.jpg"
    else:
        # 其他情况
        original_filename = f"original_frame_vehicles{vehicle_count}_{timestamp}.jpg"

    original_path = os.path.join(ORIGINAL_FOLDER, original_filename)
    cv2.imwrite(original_path, image)
    print(f"📷 原始图片已保存 (车辆数: {vehicle_count}): {original_path}")

    return original_path

# -----------------------------
# 5. 保存车辆掩码图片函数
# -----------------------------
def save_vehicle_mask(image, vehicle_bbox, mask, class_id, confidence, vehicle_index, frame_info=""):
    """
    保存完整的车辆掩码图片
    """
    if not SAVE_MASKS or mask is None:
        return None

    try:
        # 创建原图副本
        mask_img = image.copy()
        orig_h, orig_w = image.shape[:2]

        # 确保掩码是二值化的
        if mask.dtype != np.bool_ and mask.dtype != np.uint8:
            mask = mask.astype(np.float32)

        # 将掩码转换为二值图像
        if mask.max() <= 1.0:
            mask_binary = (mask > 0.5).astype(np.uint8)
        else:
            mask_binary = (mask > 127).astype(np.uint8)

        # 检查掩码尺寸，如果不匹配则调整
        if mask_binary.shape[0] != orig_h or mask_binary.shape[1] != orig_w:
            print(f"    ⚠️ 掩码尺寸不匹配: {mask_binary.shape} -> 调整到: ({orig_h}, {orig_w})")
            mask_binary = cv2.resize(mask_binary, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)

        # 创建彩色掩码
        vehicle_color = VEHICLE_COLORS.get(class_id, (0, 255, 0))  # 默认绿色
        color_mask = np.zeros_like(mask_img)

        # 应用颜色到整个掩码区域
        for c in range(3):
            color_mask[:, :, c] = np.where(mask_binary > 0, vehicle_color[c], 0)

        # 融合原图和彩色掩码（提高掩码可见性）
        mask_img = cv2.addWeighted(mask_img, 1 - MASK_ALPHA, color_mask, MASK_ALPHA, 0)

        # 绘制车辆边界框
        x_min, y_min, x_max, y_max = vehicle_bbox
        cv2.rectangle(mask_img, (x_min, y_min), (x_max, y_max), vehicle_color, 3)

        # 绘制更粗的边界框以突出显示
        cv2.rectangle(mask_img, (x_min, y_min), (x_max, y_max), vehicle_color, 2)

        # 添加车辆信息标签
        vehicle_type = VEHICLE_CLASSES.get(class_id, 'unknown')
        label = f"{vehicle_type} {confidence:.2f}"
        label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)[0]

        # 标签背景
        cv2.rectangle(mask_img, (x_min, y_min - label_size[1] - 10),
                      (x_min + label_size[0] + 10, y_min), vehicle_color, -1)
        cv2.rectangle(mask_img, (x_min, y_min - label_size[1] - 10),
                      (x_min + label_size[0] + 10, y_min), vehicle_color, 2)

        # 标签文字
        cv2.putText(mask_img, label, (x_min + 5, y_min - 5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)

        # 添加统计信息
        mask_pixels = np.sum(mask_binary)
        total_pixels = orig_h * orig_w
        mask_coverage = (mask_pixels / total_pixels) * 100

        info_lines = [
            f"Vehicle {vehicle_index} - {vehicle_type}",
            f"Confidence: {confidence:.3f}",
            f"Mask Coverage: {mask_coverage:.2f}%",
            f"Mask Pixels: {mask_pixels}",
            f"BBox: [{x_min}, {y_min}, {x_max}, {y_max}]"
        ]

        # 绘制信息背景
        text_bg_height = len(info_lines) * 25 + 20
        cv2.rectangle(mask_img, (10, 10), (400, text_bg_height), (0, 0, 0), -1)
        cv2.rectangle(mask_img, (10, 10), (400, text_bg_height), (255, 255, 255), 1)

        # 添加信息文本
        for i, line in enumerate(info_lines):
            y_pos = 35 + i * 25
            color = (255, 255, 255)
            cv2.putText(mask_img, line, (20, y_pos),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

        # 保存掩码图片
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        base_name = os.path.splitext(frame_info)[0] if frame_info else "frame"
        mask_filename = f"mask_{base_name}_vehicle{vehicle_index}_{timestamp}.jpg"
        mask_path = os.path.join(MASK_FOLDER, mask_filename)

        cv2.imwrite(mask_path, mask_img)
        print(f"📸 车辆掩码已保存 (覆盖率: {mask_coverage:.2f}%): {mask_path}")

        return mask_path

    except Exception as e:
        print(f"❌ 保存车辆掩码时出错: {e}")
        return None

# -----------------------------
# 6. 单帧处理函数（仅车辆实例分割）
# -----------------------------
def process_single_frame(frame, frame_info=""):
    """
    处理单帧图像，仅进行车辆实例分割
    返回处理结果和统计信息
    """
    if frame is None:
        return None, {}

    orig_h, orig_w = frame.shape[:2]
    print(f"处理帧: {frame_info}, 尺寸: {orig_w}x{orig_h}")

    # 使用YOLOv8进行车辆检测和实例分割
    detections = detect_vehicles_yolov8(frame)
    vehicle_count = len(detections)

    if vehicle_count == 0:
        print("    ⚠️ 未检测到车辆，跳过此帧")
        stats = {
            'total_vehicles': 0,
            'processed_vehicles': 0
        }
        return frame, stats

    print(f"    ✅ 检测到 {vehicle_count} 辆车，开始处理...")

    # 统计信息
    processed_vehicles = 0

    # 处理每个检测到的车辆
    for det_idx, det in enumerate(detections):
        x_min, y_min, x_max, y_max, conf, class_id, mask = det

        # 只处理车辆类别
        if class_id in VEHICLE_CLASSES:
            vehicle_bbox = [x_min, y_min, x_max, y_max]
            vehicle_type = VEHICLE_CLASSES.get(class_id, 'unknown')

            print(f"    🚗 车辆 {det_idx + 1}: {vehicle_type} (置信度: {conf:.2f})")

            # 保存车辆掩码图片
            if SAVE_MASKS and mask is not None:
                mask_path = save_vehicle_mask(
                    frame, vehicle_bbox, mask, class_id, conf, det_idx + 1, frame_info
                )
                if mask_path:
                    processed_vehicles += 1
                    print(f"    ✅ 车辆 {det_idx + 1} 掩码已保存")

    # 只有在检测到车辆时才保存原始图片
    if vehicle_count > 0:
        original_path = save_original_image(frame, frame_info, vehicle_count)

    # 统计信息
    stats = {
        'total_vehicles': vehicle_count,
        'processed_vehicles': processed_vehicles
    }

    print(f"    📈 本帧统计: 检测车辆数={vehicle_count}, 处理车辆数={processed_vehicles}")

    return frame, stats

# -----------------------------
# 7. RTSP流处理器
# -----------------------------
class RTSPProcessor:
    """
    RTSP流处理器，支持多路RTSP流和自动重连
    """

    def __init__(self, rtsp_url, detection_interval=3):
        self.rtsp_url = rtsp_url
        self.detection_interval = detection_interval
        self.cap = None
        self.frame_count = 0
        self.processed_frame_count = 0
        self.stats = []
        self.start_time = None
        self.running = False
        self.reconnect_attempts = 0
        self.max_reconnect_attempts = 10
        self.consecutive_error_count = 0
        self.max_consecutive_errors = 50

    def connect(self):
        """连接到RTSP流"""
        print(f"📡 尝试连接到RTSP流: {self.rtsp_url}")

        try:
            # 设置OpenCV参数以优化RTSP连接
            self.cap = cv2.VideoCapture(self.rtsp_url)

            # 设置缓冲区大小，减少延迟
            self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)

            # 设置超时参数
            self.cap.set(cv2.CAP_PROP_OPEN_TIMEOUT_MSEC, RTSP_TIMEOUT * 1000)
            self.cap.set(cv2.CAP_PROP_READ_TIMEOUT_MSEC, RTSP_TIMEOUT * 1000)

            # 检查连接是否成功
            if not self.cap.isOpened():
                print(f"❌ 无法连接到RTSP流: {self.rtsp_url}")
                return False

            # 获取视频属性
            width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            fps = self.cap.get(cv2.CAP_PROP_FPS)
            print(f"✅ 成功连接到RTSP流: {self.rtsp_url}")
            print(f"视频属性: {width}x{height}, FPS: {fps:.2f}")

            return True

        except Exception as e:
            print(f"❌ 连接RTSP流时出错: {e}")
            return False

    def reconnect(self):
        """重新连接到RTSP流"""
        print(f"🔄 尝试重新连接RTSP流: {self.rtsp_url}")

        if self.cap:
            self.cap.release()

        # 增加重连尝试次数
        self.reconnect_attempts += 1
        if self.reconnect_attempts > self.max_reconnect_attempts:
            print(f"❌ 重连尝试次数超过限制: {self.max_reconnect_attempts}")
            return False

        # 等待一段时间再重连
        time.sleep(RTSP_RECONNECT_INTERVAL)

        return self.connect()

    def process_stream(self, max_frames=None):
        """处理RTSP流"""
        if not self.cap:
            print("❌ RTSP流未连接")
            return

        print(f"🎬 开始处理RTSP流: {self.rtsp_url}")
        print(f"⚠️ 每{self.detection_interval}帧检测一次")
        print("⚠️ 按 'q' 键退出实时显示窗口")
        print("⚠️ 按 Ctrl+C 退出程序")

        self.start_time = time.time()
        self.running = True

        # 创建显示窗口
        cv2.namedWindow('RTSP Stream - Vehicle Segmentation', cv2.WINDOW_NORMAL)
        cv2.resizeWindow('RTSP Stream - Vehicle Segmentation', 800, 600)

        while self.running:
            try:
                # 读取帧
                ret, frame = self.cap.read()

                if not ret or frame is None or frame.size == 0:
                    print("⚠️ 无法读取帧，可能流已断开...")
                    self.consecutive_error_count += 1

                    if self.consecutive_error_count >= self.max_consecutive_errors:
                        print("⚠️ 连续错误次数过多，尝试重新连接...")
                        if not self.reconnect():
                            print("❌ 重新连接失败，退出处理")
                            break
                        self.consecutive_error_count = 0  # 重置错误计数
                    else:
                        time.sleep(0.1)  # 短暂等待
                        continue
                else:
                    self.consecutive_error_count = 0  # 重置错误计数
                    self.reconnect_attempts = 0  # 重置重连计数

                self.frame_count += 1

                # 显示当前帧（不包含检测结果）
                display_frame = frame.copy()
                cv2.putText(display_frame, f"Frame: {self.frame_count}", (10, 30),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
                cv2.putText(display_frame, f"RTSP: {os.path.basename(self.rtsp_url)}", (10, 60),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
                cv2.imshow('RTSP Stream - Vehicle Segmentation', display_frame)

                # 每隔指定帧数检测一次
                if self.frame_count % self.detection_interval == 0:
                    self.processed_frame_count += 1
                    print(f"--- 处理RTSP帧 {self.frame_count} (检测 {self.processed_frame_count}) ---")

                    # 处理当前帧
                    result, stats = process_single_frame(frame, f"RTSP_Frame_{self.frame_count}")

                    if result is not None:
                        stats['frame_num'] = self.frame_count
                        stats['source'] = 'RTSP'
                        self.stats.append(stats)

                        # 显示处理结果
                        cv2.putText(result, f"Processed: {self.processed_frame_count}", (10, 90),
                                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
                        cv2.imshow('RTSP Stream - Vehicle Segmentation', result)

                # 检查退出键
                key = cv2.waitKey(1) & 0xFF
                if key == ord('q'):
                    print("⏹️ 用户中断处理")
                    break

                # 限制处理帧数（如果设置了max_frames）
                if max_frames and self.frame_count >= max_frames:
                    print(f"✅ 达到最大帧数限制: {max_frames}")
                    break

            except KeyboardInterrupt:
                print("⏹️ 用户中断处理")
                break
            except Exception as e:
                print(f"⚠️ 处理RTSP帧时出错: {e}")
                self.consecutive_error_count += 1
                if self.consecutive_error_count >= self.max_consecutive_errors:
                    print("❌ 错误次数过多，尝试重新连接")
                    if not self.reconnect():
                        print("❌ 重新连接失败，退出处理")
                        break
                    self.consecutive_error_count = 0
                time.sleep(1)  # 短暂等待后继续
                continue

        # 计算处理时间
        if self.start_time:
            total_time = time.time() - self.start_time
            processing_fps = self.processed_frame_count / total_time if total_time > 0 else 0
            print(f"\n✅ RTSP流处理完成!")
            print(f"总帧数: {self.frame_count}, 处理帧数: {self.processed_frame_count}")
            print(f"总耗时: {total_time:.2f}秒, 处理速度: {processing_fps:.2f} FPS")

        # 关闭显示窗口
        cv2.destroyWindow('RTSP Stream - Vehicle Segmentation')

    def release(self):
        """释放资源"""
        self.running = False
        if self.cap:
            self.cap.release()
        cv2.destroyAllWindows()

# -----------------------------
# 8. RTSP流处理函数
# -----------------------------
def process_rtsp_streams():
    """
    处理RTSP流
    """
    if not RTSP_URLS:
        print("❌ 没有配置RTSP URL！请在RTSP_URLS列表中添加RTSP流地址")
        return

    print(f"📡 检测到 {len(RTSP_URLS)} 个RTSP流")

    for i, rtsp_url in enumerate(RTSP_URLS, 1):
        print(f"\n--- 处理RTSP流 {i}/{len(RTSP_URLS)}: {rtsp_url} ---")

        # 创建RTSP处理器
        rtsp_processor = RTSPProcessor(rtsp_url, DETECTION_INTERVAL)

        # 连接并处理流
        if rtsp_processor.connect():
            rtsp_processor.process_stream()

            # 输出统计报告
            print_statistics(rtsp_processor.stats, f"RTSP流 {i}")
        else:
            print(f"❌ 无法连接到RTSP流 {i}: {rtsp_url}")

        # 释放资源
        rtsp_processor.release()

# -----------------------------
# 9. 统计报告函数
# -----------------------------
def print_statistics(stats, data_type):
    """
    输出统计报告
    """
    if not stats:
        print(f"❌ 没有{data_type}处理数据！")
        return

    print(f"\n{'=' * 50}")
    print(f"🚗 {data_type}车辆实例分割统计报告")
    print(f"{'=' * 50}")

    total_frames = len(stats)
    total_vehicles = sum(stat['total_vehicles'] for stat in stats)
    total_processed = sum(stat['processed_vehicles'] for stat in stats)

    print(f"总处理帧数: {total_frames}")
    print(f"总检测车辆数: {total_vehicles}")
    print(f"总处理车辆数: {total_processed}")
    print(f"检测间隔: 每{DETECTION_INTERVAL}帧检测一次")

    # 输出保存信息
    if SAVE_MASKS:
        print(f"车辆掩码保存: 已启用 (保存在 {MASK_FOLDER} 文件夹)")

    if SAVE_ORIGINAL:
        print(f"原始图片保存: 已启用 (仅在检测到车辆时保存，保存在 {ORIGINAL_FOLDER} 文件夹)")

# -----------------------------
# 10. 主函数
# -----------------------------
def main():
    """
    主函数 - 仅处理RTSP流
    """
    print("=" * 60)
    print("🚗 车辆实例分割系统（基于YOLOv8n-seg）")
    print("=" * 60)
    print("配置信息:")
    print(f"  - 车辆检测模型: YOLOv8n-seg")
    print(f"  - 检测间隔: 每{DETECTION_INTERVAL}帧检测一次")
    print(f"  - 掩码保存: {'启用' if SAVE_MASKS else '禁用'}")
    print(f"  - 原始图片: {'启用' if SAVE_ORIGINAL else '禁用'} (仅在检测到车辆时)")
    print(f"  - 掩码透明度: {MASK_ALPHA}")
    print()

    print("🎯 开始处理RTSP流...")
    process_rtsp_streams()

    print(f"\n🎉 RTSP流处理完成！")
    print(f"📁 结果目录:")
    if SAVE_MASKS:
        print(f"   - 车辆掩码: {MASK_FOLDER}")
    if SAVE_ORIGINAL:
        print(f"   - 原始图片: {ORIGINAL_FOLDER} (仅在检测到车辆时保存)")

# -----------------------------
# 11. 运行主函数
# -----------------------------
if __name__ == "__main__":
    main()